/*******************************************************************************
 * Copyright (c) 2014 Steffen Mikkelsen
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the GNU Public License v3.0
 * which accompanies this distribution, and is available at
 * http://www.gnu.org/licenses/gpl.html
 * 
 * Contributors:
 *     Steffen Mikkelsen - initial API and implementation
 ******************************************************************************/
package com.subsecret.secureconnection.securestream.encryption;

import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.X509EncodedKeySpec;

import javax.crypto.KeyAgreement;
import javax.crypto.interfaces.DHPrivateKey;
import javax.crypto.interfaces.DHPublicKey;

/**
 * Diffie Hellman implementation for Server-Client communication
 * 
 * @author steffen
 * 
 */
public class DiffieHellmanServer {
    
    private DHPrivateKey senderPrivate = null;
    private byte[] secret = null;
    
    public DiffieHellmanServer() {
    }
    
    /**
     * Generate public key to be transferred to Client
     * 
     * @return byte[] to be given to client
     * @throws NoSuchAlgorithmException
     */
    public byte[] generateServerPart1() {
        
        KeyPairGenerator kpg;
        try {
            kpg = KeyPairGenerator.getInstance("DiffieHellman");
            kpg.initialize(1024);
            
            KeyPair senderDHKeyPair = kpg.generateKeyPair();
            DHPublicKey senderPublic = (DHPublicKey) senderDHKeyPair.getPublic();
            senderPrivate = (DHPrivateKey) senderDHKeyPair.getPrivate();
            
            // send public code (both generator and public key)
            byte[] encodedPublic = senderPublic.getEncoded();
            return encodedPublic;
        }
        catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e.getClass().getSimpleName() + ": " + e.getMessage());
        }        
        
    }
    
    /**
     * Finish initialization of Server part with public key from client
     * 
     * @param receivedPublic
     *            Public key from Client
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeySpecException
     * @throws InvalidKeyException
     */
    public void generateServerPart2(byte[] receivedPublic)  {
        X509EncodedKeySpec x509Key = new X509EncodedKeySpec(receivedPublic);
        
        try {
            KeyFactory kfactory = KeyFactory.getInstance("DiffieHellman");
            DHPublicKey receivedPublicKey = (DHPublicKey) kfactory.generatePublic(x509Key);
            KeyAgreement senderAgree = KeyAgreement.getInstance("DiffieHellman");
            senderAgree.init(senderPrivate);
            senderAgree.doPhase(receivedPublicKey, true);
            secret = senderAgree.generateSecret();
        }
        catch (InvalidKeyException | IllegalStateException | NoSuchAlgorithmException | InvalidKeySpecException e) {
            throw new RuntimeException(e.getClass().getSimpleName() + ": " + e.getMessage());
        }
    }
    
    /**
     * @return The shared secret, available after part 2 has finished
     */
    public byte[] getSharedSecret() {
        return secret;
    }
    
}
